from typing import Annotated, Any, Union

from typing_inspect import get_args, get_origin, is_generic_type, is_literal_type, is_union_type

from zkl_pyutils_registries.basics import Matcher, NotMatched
from .type_matcher import match_type


class SubTypehintMatcher(Matcher[Any, Any, tuple]):
    """ Matches if the key typehint is subset of the query typehint """

    def match(self, query: Any, key: Any) -> tuple:
        return match_typehint(key, query)


class SupTypehintMatcher(Matcher[Any, Any, tuple]):
    """ Matches if the key typehint is superset of the query typehint """

    def match(self, query: Any, key: Any) -> tuple:
        return match_typehint(query, key)


def match_typehint(sub_typehint: Any, sup_typehint: Any):
    sub_typehint = type(None) if sub_typehint is None else sub_typehint
    sup_typehint = type(None) if sup_typehint is None else sup_typehint
    return match_typehint_equal(sub_typehint, sup_typehint)


def match_typehint_equal(sub_typehint: Any, sup_typehint: Any) -> tuple:
    if sub_typehint is sup_typehint:
        return +2,
    if sub_typehint == sup_typehint:
        return +1,
    return 0, *match_typehint_annotated(sub_typehint, sup_typehint)


def match_typehint_annotated(sub_typehint: Any, sup_typehint: Any) -> tuple:
    from typing import get_origin
    sub_typehint_is_annotated = get_origin(sub_typehint) is Annotated
    sup_typehint_is_annotated = get_origin(sup_typehint) is Annotated
    sub_typehint_origin = sub_typehint.__origin__ if sub_typehint_is_annotated else sub_typehint
    sup_typehint_origin = sup_typehint.__origin__ if sup_typehint_is_annotated else sup_typehint
    return match_typehint_union(sub_typehint_origin, sup_typehint_origin)


def match_typehint_union(sub_typehint: Any, sup_typehint: Any) -> tuple:
    if not is_union_type(sub_typehint) and not is_union_type(sup_typehint):
        return 0, *match_typehint_any(sub_typehint, sup_typehint)

    sup_typehint_args = get_args(sup_typehint) if is_union_type(sup_typehint) else (sup_typehint,)
    sub_typehint_args = get_args(sub_typehint) if is_union_type(sub_typehint) else (sub_typehint,)

    worst_match = None
    sup_redundancy = [True] * len(sup_typehint_args)
    for sub_typehint_arg in sub_typehint_args:
        best_match = None
        for i, sup_typehint_arg in enumerate(sup_typehint_args):
            try:
                match = match_typehint(sub_typehint_arg, sup_typehint_arg)
                sup_redundancy[i] = False
                if best_match is None or match > best_match:
                    best_match = match
            except NotMatched:
                continue
        if best_match is None:
            raise NotMatched
        if worst_match is None or best_match < worst_match:
            worst_match = best_match
    if worst_match is None:
        raise NotMatched

    return -1, -sum(sup_redundancy), worst_match


def match_typehint_any(sub_typehint: Any, sup_typehint: Any) -> tuple:
    if sup_typehint is Any:
        if sub_typehint is Any:
            return +1,
        else:
            return -1,

    if sub_typehint is Any:
        raise NotMatched

    return 0, *match_typehint_literal(sub_typehint, sup_typehint)


def match_typehint_literal(sub_typehint: Any, sup_typehint: Any) -> tuple:
    if is_literal_type(sub_typehint):
        sub_typehint_args = set(get_args(sub_typehint))
        if is_literal_type(sup_typehint):
            sup_typehint_args = set(get_args(sup_typehint))
            if not sub_typehint_args.issubset(sup_typehint_args):
                raise NotMatched
            return +1, -len(sup_typehint_args - sub_typehint_args)

        involved_types = set()
        for sub_typehint_arg in sub_typehint_args:
            sub_typehint_arg_type = type(sub_typehint_arg)
            involved_types.add(sub_typehint_arg_type)
        return -1, match_typehint(Union[tuple(involved_types)], sup_typehint)

    if is_literal_type(sup_typehint):
        raise NotMatched

    return 0, *match_typehint_type(sub_typehint, sup_typehint)


def match_typehint_type(sub_typehint: Any, sup_typehint: Any) -> tuple:
    sub_type = get_origin(sub_typehint) or sub_typehint
    if not isinstance(sub_type, type):
        raise TypeError(f"Cannot deal with {sub_type}!")

    sup_type = get_origin(sup_typehint) or sup_typehint
    if not isinstance(sup_type, type):
        raise TypeError(f"Cannot deal with {sup_type}!")

    return match_type(sub_type, sup_type), *match_typehint_generic(sub_typehint, sup_typehint)


def match_typehint_generic(sub_typehint: Any, sup_typehint: Any) -> tuple:
    sub_type = get_origin(sub_typehint) or sub_typehint
    sup_type = get_origin(sup_typehint) or sup_typehint
    if not sub_type == sup_type:
        return ()

    sub_typehint_args = get_args(sub_typehint) if is_generic_type(sub_typehint) else ()
    sup_typehint_args = get_args(sup_typehint) if is_generic_type(sup_typehint) else ()

    len_match = -abs(len(sub_typehint_args) - len(sup_typehint_args))
    args_match = tuple(
        match_typehint(sub_typehint_arg, sup_typehint_arg)
        for sub_typehint_arg, sup_typehint_arg in zip(sub_typehint_args, sup_typehint_args))
    return len_match, args_match
